import sqlite3
import os
import argparse
import shutil
from datetime import datetime
import hashlib
import random
import string

def create_database(db_path):
    """创建包和文件表结构"""
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    # 创建文件表
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS files (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        file_name TEXT NOT NULL,
        relative_path TEXT NOT NULL,
        file_type TEXT,
        file_size INTEGER,
        content BLOB NOT NULL,
        created_at DATETIME NOT NULL, 
        block_type TEXT DEFAULT "file"
    )
    ''')
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS metadata (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        info_head TEXT NOT NULL UNIQUE,
        info_body TEXT NOT NULL
    )
                   ''')
    # 创建passkey表
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS passkey (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        key_hash TEXT NOT NULL,
        created_at DATETIME NOT NULL
    )
    ''')
    conn.commit()
    conn.close()

def xor_encrypt(content, key):
    """使用异或进行加密/解密"""
    key_bytes = key.encode()
    key_length = len(key_bytes)
    result = bytearray()
    
    for i in range(len(content)):
        result.append(content[i] ^ key_bytes[i % key_length])
    
    return bytes(result)

def get_or_create_passkey(conn, password=None):
    """获取或创建密码密钥"""
    # 如果没有提供密码，返回空字符串，表示不加密
    if password is None:
        return ""
    
    cursor = conn.cursor()
    
    # 检查是否已有密码
    cursor.execute("SELECT key_hash FROM passkey ORDER BY id DESC LIMIT 1")
    existing_key = cursor.fetchone()
    
    if existing_key:
        return existing_key[0]
    
    # 生成MD5哈希作为密钥
    key_hash = hashlib.md5(password.encode()).hexdigest()
    
    # 保存密钥哈希
    cursor.execute("INSERT INTO passkey (key_hash, created_at) VALUES (?, ?)", 
                 (key_hash, datetime.now()))
    conn.commit()
    
    return key_hash

def insert_file(conn, file_path, relative_path, password=None):
    """读取文件并插入包（支持加密）"""
    # 获取文件信息
    file_name = os.path.basename(file_path)
    file_type = os.path.splitext(file_name)[1].lower() if os.path.splitext(file_name)[1] else "unknown"
    file_size = os.path.getsize(file_path)
    
    # 获取密钥
    key_hash = get_or_create_passkey(conn, password)
    
    # 读取文件内容
    try:
        with open(file_path, 'rb') as f:
            content = f.read()
        
        # 如果有密钥则加密内容，否则不加密
        if key_hash:
            final_content = xor_encrypt(content, key_hash)
        else:
            final_content = content
        
        # 插入包
        cursor = conn.cursor()
        cursor.execute('''
        INSERT INTO files (file_name, relative_path, file_type, file_size, content, created_at)
        VALUES (?, ?, ?, ?, ?, ?)
        ''', (file_name, relative_path, file_type, file_size, final_content, datetime.now()))
        
        conn.commit()
        print(f"成功添加文件: {relative_path}")
        return True
    except Exception as e:
        print(f"处理文件 {file_path} 时出错: {str(e)}")
        return False
def write_metadata(conn, info_head, info_body):
    """写入元数据"""
    cursor = conn.cursor()
    cursor.execute('''
    INSERT INTO metadata (info_head, info_body)
    VALUES (?, ?)
    ''', (info_head, info_body))
    conn.commit()
def get_metadata(conn, info_head):
    """获取元数据"""
    cursor = conn.cursor()
    cursor.execute('''
    SELECT info_head, info_body FROM metadata WHERE info_head=?
    ''', (info_head,))
    metadata = cursor.fetchone()
    return metadata

def insert_placeholder_file(conn, folder_path, root_dir, password=None):
    """在空文件夹中插入.placeholder文件"""
    # 创建一个临时的placeholder文件
    import tempfile
    with tempfile.NamedTemporaryFile(delete=False, mode='wb') as temp:
        temp.write(b'')  # 空内容
        temp_path = temp.name
    
    try:
        # 计算相对路径
        relative_folder = os.path.relpath(folder_path, root_dir)
        placeholder_relative_path = os.path.join(relative_folder, '.placeholder')
        
        # 插入placeholder文件
        insert_file(conn, temp_path, placeholder_relative_path, password)
        print(f"在空文件夹 {relative_folder} 中创建了.placeholder文件")
    finally:
        # 删除临时文件
        if os.path.exists(temp_path):
            os.remove(temp_path)

def pack_files(source_path, db_path, password=None):
    """打包文件/文件夹到包（支持加密）"""
    # 创建包
    create_database(db_path)
    conn = sqlite3.connect(db_path)
    write_metadata(conn, "package_data", "block.enabled=1")
    write_metadata(conn, "is_a_package", "1")
    try:
        if os.path.isfile(source_path):
            # 处理单个文件
            relative_path = os.path.basename(source_path)
            insert_file(conn, source_path, relative_path, password)
        elif os.path.isdir(source_path):
            # 处理文件夹
            root_dir = os.path.abspath(source_path)
            
            # 先处理所有文件
            for root, dirs, files in os.walk(source_path):
                for file in files:
                    file_path = os.path.join(root, file)
                    relative_path = os.path.relpath(file_path, root_dir)
                    insert_file(conn, file_path, relative_path, password)
            
            # 再检查所有空文件夹并创建.placeholder文件
            for root, dirs, files in os.walk(source_path):
                # 检查是否是空文件夹（不包含任何文件）
                if not files and os.listdir(root) == dirs:
                    insert_placeholder_file(conn, root, root_dir, password)
        else:
            print(f"错误: {source_path} 不是有效的文件或文件夹")
            return False
            
        print(f"打包完成，包文件已保存至: {db_path}")
        return True
    except Exception as e:
        print(f"打包过程中出错: {str(e)}")
        return False
    finally:
        conn.close()

def unpack_files(db_path, dest_path, password=None):
    """从包解包文件（支持解密）"""
    if not os.path.exists(db_path):
        print(f"错误: 包文件 {db_path} 不存在")
        return False
    
    try:
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        if not get_metadata(conn, "is_a_package"):
            print("错误: 指定的文件不是有效的包文件")
            return False
        
        # 获取密钥
        cursor.execute("SELECT key_hash FROM passkey ORDER BY id DESC LIMIT 1")
        saved_key = cursor.fetchone()
        
        # 如果有保存的密钥，则需要密码解密
        if saved_key:
            if password is None:
                print("错误: 该包已加密，需要提供密码 (-p 参数)")
                return False
            
            # 验证密码
            provided_key_hash = hashlib.md5(password.encode()).hexdigest()
            if provided_key_hash != saved_key[0]:
                print("错误: 密码不正确")
                return False
            
            key_hash = provided_key_hash
        else:
            # 如果没有加密，key_hash为空
            key_hash = ""
        
        # 获取所有文件
        cursor.execute("SELECT file_name, relative_path, content FROM files")
        files = cursor.fetchall()
        
        if not files:
            print("包中没有找到文件")
            return True
        
        # 创建目标目录
        os.makedirs(dest_path, exist_ok=True)
        
        # 解包每个文件
        for file in files:
            file_name, relative_path, encrypted_content = file
            file_path = os.path.join(dest_path, relative_path)
            
            # 创建目录
            os.makedirs(os.path.dirname(file_path), exist_ok=True)
            
            # 解密内容（如果有密钥）
            if key_hash:
                content = xor_encrypt(encrypted_content, key_hash)
            else:
                content = encrypted_content
            
            # 写入文件
            with open(file_path, 'wb') as f:
                f.write(content)
            
            print(f"解包文件: {relative_path}")
        
        print(f"解包完成，文件已保存至: {dest_path}")
        return True
    except Exception as e:
        print(f"解包过程中出错: {str(e)}")
        return False
    finally:
        if 'conn' in locals():
            conn.close()

def main():
    parser = argparse.ArgumentParser(description='文件打包解包工具 - 将文件/文件夹打包到SQLite包或从包解包文件')
    subparsers = parser.add_subparsers(dest='command', help='命令')
    
    # 打包命令
    pack_parser = subparsers.add_parser('pack', help='打包文件/文件夹到包')
    pack_parser.add_argument('-s', '--source', required=True, help='要打包的文件或文件夹路径')
    pack_parser.add_argument('-d', '--dest', required=True, help='输出的包文件路径')
    pack_parser.add_argument('-p', '--password', help='加密密码，如果不提供则随机生成')
    
    # 解包命令
    unpack_parser = subparsers.add_parser('unpack', help='从包解包文件')
    unpack_parser.add_argument('-s', '--source', required=True, help='包文件路径')
    unpack_parser.add_argument('-d', '--dest', required=True, help='解包目标目录')
    unpack_parser.add_argument('-p', '--password', help='解密密码')
    
    args = parser.parse_args()
    
    if args.command == 'pack':
        pack_files(args.source, args.dest, args.password)
    elif args.command == 'unpack':
        unpack_files(args.source, args.dest, args.password)
    else:
        parser.print_help()

if __name__ == "__main__":
    main()
